from src.agents.retrieval_agent import RetrievalAgent
from src.utils.citation_manager import CitationManager
from src.agents.review_generator import ReviewGenerator
from src.templates.prompts import Prompts
from zhipuai import ZhipuAI
from datetime import datetime
from typing import Optional, Dict, List, Any
import re
import time
from .api_calls import search_papers
import json
import asyncio

MODEL = "glm-4-plus"
client = ZhipuAI(api_key='YOUR API KEY')

class RAG:
    def __init__(self, min_citations: int = 5):
        self.min_citations = min_citations
        self.citations = []
        self.citation_manager = CitationManager()
        self.max_papers = 50
        self.retrieval_agent = RetrievalAgent()
        self.prompts = Prompts()

    def _extract_year(self, filename: str) -> str:
        year_match = re.search(r'20\d{2}', filename)
        return year_match.group() if year_match else str(datetime.now().year)

    def _validate_paper_data(self, paper: dict):
        """验证论文数据的完整性"""
        required_fields = ['paper_title', 'venue', 'year', 'chunk_id', 'content', 'paper_id']
        missing = [f for f in required_fields if not paper.get(f)]
        if missing:
            raise ValueError(f"论文数据缺失必要字段: {', '.join(missing)}")
        return True

    def _process_paper(self, paper: dict) -> dict:
        """处理论文数据，确保数据完整性"""
        self._validate_paper_data(paper)
        return {
            'title': paper['title'],
            'venue': paper['venue'],
            'year': paper['year'],
            'chunk_id': paper['chunk_id'],
            'content': paper['content'],
            'paper_id': paper['paper_id']
        }

    async def _retrieve_papers(self, query: str, min_papers: int) -> List[Dict[str, Any]]:
        print(f"开始检索论文，查询关键词: {query}，最小论文数量: {min_papers}")
        try:
            if not hasattr(self, 'retrieval_agent'):
                self.retrieval_agent = RetrievalAgent()
            papers = self.retrieval_agent.search_papers(query, top_k=min_papers)
            if not papers:
                return []
            
            processed_papers = []
            for paper in papers:
                try:
                    if not paper.get('entity', {}).get('paper_title'):
                        continue
                    
                    processed_paper = {
                        'title': paper.get('title'),
                        'venue': paper.get('venue', 'Conference'),
                        'year': paper.get('year', str(datetime.now().year)),
                        'chunk_id': paper.get('chunk_id', str(len(processed_papers))),
                        'content': paper.get('content', ''),
                        'paper_id': paper.get('paper_id', f'unknown_{len(processed_papers)}')
                    }
                    processed_papers.append(processed_paper)
                except Exception as e:
                    print(f"处理论文数据时出错: {str(e)}")
                    continue
            
            print(f"处理完成，找到{len(processed_papers)}篇有效论文")
            return processed_papers[:self.max_papers]
            
            
        except Exception as e:
            print(f"论文检索失败: {str(e)}")
            return []

    _REVIEW_TYPE_MAP = {
        'status': 'RESEARCH_STATUS_REVIEW',
        'comparative': 'COMPARATIVE_ANALYSIS',
        'trajectory': 'METHOD_TRAJECTORY',
        'concept': 'TECHNICAL_CONCEPT_REVIEW'
    }

    async def _generate_content(self, papers: List[Dict[str, Any]], review_type: str, topic: str) -> Optional[str]:
        print("开始生成综述内容...")
        try:
            # 重置并初始化引用管理器
            self.citation_manager.reset()
            self.citation_manager.set_topic(topic)
            
            # 优化论文信息的展示格式
            papers_text = "\n".join([f"{i+1}. {p['title']} ({p['venue']}, {p['year']})\nID: {p['paper_id']}, Chunk: {p['chunk_id']}\n{p['content']}\n" 
                                   for i, p in enumerate(papers)])
            
            # 合并所有模板内容
            combined_prompt = """
            {base_requirements} 
            {technical_concept}
            {research_status}
            {comparative_analysis}
            {method_trajectory}
            """.format(
                base_requirements=self.prompts.BASE_REQUIREMENTS,
                topic=topic,
                technical_concept=self.prompts.TECHNICAL_CONCEPT_REVIEW,
                research_status=self.prompts.RESEARCH_STATUS_REVIEW,
                comparative_analysis=self.prompts.COMPARATIVE_ANALYSIS,
                method_trajectory=self.prompts.METHOD_TRAJECTORY
            )
            
            print("正在调用API生成内容...")
            response = client.chat.completions.create(
                model=MODEL,
                messages=[{
                    "role": "system", 
                    "content": "You are an expert academic writer. Your task is to review and enhance the readability of the provided  text in a research paper. Ensure that the text is clear, concise, and free from jargon while maintaining its academic integrity. Make sure to include all citations as specified in the template."
                }, 
                {"role": "user", "content": combined_prompt}],
                temperature=0.1,
                max_tokens=4000,
            )
            
            content = response.choices[0].message.content
            print("API调用成功，开始处理引用...")
            
            # 处理引用标记，转换为<sup>number</sup>格式
            content, citations = self.citation_manager.process_citations(content, papers)
            
            # 将[n]格式转换为<sup>n</sup>格式
            content = re.sub(r'\[(\d+)\]', r'<sup>\1</sup>', content)
            

            # 更新实例的引用列表
            self.citations = self.citation_manager.get_citations()
            print(f"内容生成完成，包含{len(self.citations)}个引用")
            return content
            
        except Exception as e:
            print(f"生成内容失败: {str(e)}")
            return None
            
        except Exception as e:
            print(f"生成内容失败: {str(e)}")
            return None

    async def generate_review(
            self, 
            query: str, 
            review_type: str = 'status',
            min_papers: int = 50
        ) -> Optional[Dict[str, Any]]:
        self.citations = []

        papers = await self._retrieve_papers(query, min_papers)
        if not papers:
            papers = await self._retrieve_papers(query, max(10, min_papers // 2))
            if not papers:
                return None

        content = await self._generate_content(papers, review_type, query)
        if not content:
            return None

        return {
            'content': content,
            'citations': self.citations,
            'papers': papers
        }